package sync

import (
	"context"
	"sync"

	"sieve_engine/lib/app"
)

func NewCallGroup[T any](ctx context.Context) *CallGroup[T] {
	return &CallGroup[T]{
		ctx: ctx,
		fns: []func(context.Context) (T, error){},
	}
}

type CallGroup[T any] struct {
	ctx    context.Context
	fns    []func(context.Context) (T, error)
	result []*CallResult[T]
	wg     sync.WaitGroup
}

type CallResult[T any] struct {
	Data T
	Err  error
}

func (g *CallGroup[T]) Run(fn func(context.Context) (T, error)) {
	g.fns = append(g.fns, fn)
}

func (g *CallGroup[T]) Wait() []*CallResult[T] {
	g.result = make([]*CallResult[T], len(g.fns))
	for i := range g.fns {
		g.run(i)
	}
	g.wg.Wait()
	return g.result
}

func (g *CallGroup[T]) run(i int) {
	g.wg.Add(1)
	app.Async(func() {
		defer g.wg.Done()
		ch := g.runFn(i)
		select {
		case <-g.ctx.Done():
			g.result[i] = &CallResult[T]{Err: g.ctx.Err()}
			return
		case r := <-ch:
			g.result[i] = r
			return
		}
	})
}

func (g *CallGroup[T]) runFn(i int) chan *CallResult[T] {
	ch := make(chan *CallResult[T])
	app.Async(func() {
		res, err := g.fns[i](g.ctx)
		select {
		case ch <- &CallResult[T]{res, err}:
			break
		default:
			break
		}
	})
	return ch
}
